# coding=utf-8
"""
USER: Klaus
DATE: 2018/1/10
TODO: Katana_component
"""
import maya.cmds as cmds
import pymel.core as pm


def list_dictionary():
    trans_sel = pm.ls(typ='transform')

    u"""角色部分"""
    sel_chr = [i for i in trans_sel if pm.hasAttr(i, 'asset_name') and 'c' in i.name().split('_')[0]]
    chr_attr_list = []
    for x in sel_chr:
        for i in x.getChildren():
            if 'geometry' in i.name():
                chr_attr_list.append(i)

    u"""添加 assembly 与 component 属性"""
    for i in chr_attr_list:
        for n in i.getChildren():
            if not pm.hasAttr(n, 'sgxml_nodeGroupType'):
                cmds.addAttr(str(n.longName()), ln='sgxml_nodeGroupType', dt='string')
                cmds.setAttr(n + '.sgxml_nodeGroupType', 'assembly', type='string')
                for x in n.getChildren():
                    if not pm.hasAttr(x, 'sgxml_nodeType'):
                        cmds.addAttr(str(x.longName()), ln='sgxml_nodeType', dt='string')
                        cmds.setAttr(x + '.sgxml_nodeType', 'component', type='string')
                    if not pm.hasAttr(x, 'sgxml_refType'):
                        cmds.addAttr(str(x.longName()), ln='sgxml_refType', dt='string')
                        cmds.setAttr(x + '.sgxml_refType', 'abc', type='string')
    assembly_chr_dict = {}
    sel_chr_assembly = [i.name() for i in pm.ls(typ='transform') if pm.hasAttr(i, 'sgxml_nodeGroupType') and 'c' in i.name().split('_')[0]]
    alembic_chr_dict = {}
    sel_chr_alembic = [i.name() for i in pm.ls(typ='transform') if pm.hasAttr(i, 'sgxml_nodeType') and 'c' in i.name().split('_')[0]]
    for x in sel_chr_assembly:
        if not assembly_chr_dict.has_key('assembly'):
            assembly_chr_dict['assembly'] = [x]
        else:
            assembly_chr_dict['assembly'].append(x)
    for x in sel_chr_alembic:
        if not alembic_chr_dict.has_key('component'):
            alembic_chr_dict['component'] = [x]
        else:
            alembic_chr_dict['component'].append(x)

    u"""拼装角色字典"""
    chr_dict = {}
    for i in chr_attr_list:
        if not chr_dict.has_key(i.name()):
            for k_ass, v_ass in assembly_chr_dict.items():
                ass_dict = {}
                ass = [x for x in v_ass if x.startswith(i.split(':')[0])]
                ass = [ass[n] for n in range(len(ass)) if ass[n] not in ass[:n]]
                ass_dict[k_ass] = ass
            for k_abc, v_abc in alembic_chr_dict.items():
                abc_dict = {}
                abc = [x for x in v_abc if x.startswith(i.split(':')[0])]
                abc = [abc[n] for n in range(len(abc)) if abc[n] not in abc[:n]]
                abc_dict[k_abc] = abc
            chr_merge_dict = dict(ass_dict.items() + abc_dict.items())
            chr_dict[i.name()] = chr_merge_dict
        else:
            for k_ass, v_ass in assembly_chr_dict.items():
                ass_dict = {}
                ass = [x for x in v_ass if x.startswith(i.split(':')[0])]
                ass = [ass[n] for n in range(len(ass)) if ass[n] not in ass[:n]]
                ass_dict[k_ass] = ass
            for k_abc, v_abc in alembic_chr_dict.items():
                abc_dict = {}
                abc = [x for x in v_abc if x.startswith(i.split(':')[0])]
                abc = [abc[n] for n in range(len(abc)) if abc[n] not in abc[:n]]
                abc_dict[k_abc] = abc
            chr_merge_dict = dict(ass_dict.items() + abc_dict.items())
            chr_dict[i.name()].append(chr_merge_dict)

    u"""场景部分"""
    sel_bg = [i for i in trans_sel if pm.hasAttr(i, 'asset_name') and 'b' in i.name().split('_')[0]]
    bg_attr_list = []
    for x in sel_bg:
        for i in x.getChildren():
            if 'geometry' in i.name():
                bg_attr_list.append(i)

    u"""添加 assembly 与 component 属性"""
    for i in bg_attr_list:
        for n in i.getChildren():
            if not pm.hasAttr(n, 'sgxml_nodeGroupType'):
                cmds.addAttr(str(n.longName()), ln='sgxml_nodeGroupType', dt='string')
                cmds.setAttr(n + '.sgxml_nodeGroupType', 'assembly', type='string')
                for x in n.getChildren():
                    if not pm.hasAttr(x, 'sgxml_nodeType'):
                        cmds.addAttr(str(x.longName()), ln='sgxml_nodeType', dt='string')
                        cmds.setAttr(x + '.sgxml_nodeType', 'component', type='string')
                    if not pm.hasAttr(x, 'sgxml_refType'):
                        cmds.addAttr(str(x.longName()), ln='sgxml_refType', dt='string')
                        cmds.setAttr(x + '.sgxml_refType', 'abc', type='string')
    assembly_bg_dict = {}
    sel_bg_assembly = [i.name() for i in pm.ls(typ='transform') if pm.hasAttr(i, 'sgxml_nodeGroupType') and 'b' in i.name().split('_')[0]]
    alembic_bg_dict = {}
    sel_bg_alembic = [i.name() for i in pm.ls(typ='transform') if pm.hasAttr(i, 'sgxml_nodeType') and 'b' in i.name().split('_')[0]]
    for x in sel_bg_assembly:
        if not assembly_bg_dict.has_key('assembly'):
            assembly_bg_dict['assembly'] = [x]
        else:
            assembly_bg_dict['assembly'].append(x)

    for x in sel_bg_alembic:
        if not alembic_bg_dict.has_key('component'):
            alembic_bg_dict['component'] = [x]
        else:
            alembic_bg_dict['component'].append(x)

    u"""拼装场景字典"""
    bg_dict = {}
    for i in bg_attr_list:
        if not bg_dict.has_key(i.name()):
            for k_ass, v_ass in assembly_bg_dict.items():
                ass_dict = {}
                ass = [x for x in v_ass if x.startswith(i.split(':')[0])]
                ass = [ass[n] for n in range(len(ass)) if ass[n] not in ass[:n]]
                ass_dict[k_ass] = ass
            for k_abc, v_abc in alembic_bg_dict.items():
                abc_dict = {}
                abc = [x for x in v_abc if x.startswith(i.split(':')[0])]
                abc = [abc[n] for n in range(len(abc)) if abc[n] not in abc[:n]]
                abc_dict[k_abc] = abc
            bg_merge_dict = dict(ass_dict.items() + abc_dict.items())
            bg_dict[i.name()] = bg_merge_dict
        else:
            for k_ass, v_ass in assembly_bg_dict.items():
                ass_dict = {}
                ass = [x for x in v_ass if x.startswith(i.split(':')[0])]
                ass = [ass[n] for n in range(len(ass)) if ass[n] not in ass[:n]]
                ass_dict[k_ass] = ass
            for k_abc, v_abc in alembic_bg_dict.items():
                abc_dict = {}
                abc = [x for x in v_abc if x.startswith(i.split(':')[0])]
                abc = [abc[n] for n in range(len(abc)) if abc[n] not in abc[:n]]
                abc_dict[k_abc] = abc
            bg_merge_dict = dict(ass_dict.items() + abc_dict.items())
            bg_dict[i.name()].append(bg_merge_dict)

    u"""道具部分"""
    sel_prop = [i for i in trans_sel if pm.hasAttr(i, 'asset_name') and 'p' in i.name().split('_')[0]]
    prop_attr_list = []
    for x in sel_prop:
        for i in x.getChildren():
            if 'geometry' in i.name():
                prop_attr_list.append(i)

    u"""添加 assembly 与 component 属性"""
    for i in prop_attr_list:
        for n in i.getChildren():
            if not pm.hasAttr(n, 'sgxml_nodeGroupType'):
                cmds.addAttr(str(n.longName()), ln='sgxml_nodeGroupType', dt='string')
                cmds.setAttr(n + '.sgxml_nodeGroupType', 'assembly', type='string')
                for x in n.getChildren():
                    if not pm.hasAttr(x, 'sgxml_nodeType'):
                        cmds.addAttr(str(x.longName()), ln='sgxml_nodeType', dt='string')
                        cmds.setAttr(x + '.sgxml_nodeType', 'component', type='string')
                    if not pm.hasAttr(x, 'sgxml_refType'):
                        cmds.addAttr(str(x.longName()), ln='sgxml_refType', dt='string')
                        cmds.setAttr(x + '.sgxml_refType', 'abc', type='string')
    assembly_prop_dict = {}
    sel_prop_assembly = [i.name() for i in pm.ls(typ='transform') if pm.hasAttr(i, 'sgxml_nodeGroupType') and 'p' in i.name().split('_')[0]]
    alembic_prop_dict = {}
    sel_prop_alembic = [i.name() for i in pm.ls(typ='transform') if pm.hasAttr(i, 'sgxml_nodeType') and 'p' in i.name().split('_')[0]]
    for x in sel_prop_assembly:
        if not assembly_prop_dict.has_key('assembly'):
            assembly_prop_dict['assembly'] = [x]
        else:
            assembly_prop_dict['assembly'].append(x)
    for x in sel_prop_alembic:
        if not alembic_prop_dict.has_key('component'):
            alembic_prop_dict['component'] = [x]
        else:
            alembic_prop_dict['component'].append(x)

    u"""拼装道具字典"""
    prop_dict = {}
    for i in prop_attr_list:
        if not prop_dict.has_key(i.name()):
            for k_ass, v_ass in assembly_prop_dict.items():
                ass_dict = {}
                ass = [x for x in v_ass if x.startswith(i.split(':')[0])]
                ass = [ass[n] for n in range(len(ass)) if ass[n] not in ass[:n]]
                ass_dict[k_ass] = ass
            for k_abc, v_abc in alembic_prop_dict.items():
                abc_dict = {}
                abc = [x for x in v_abc if x.startswith(i.split(':')[0])]
                abc = [abc[n] for n in range(len(abc)) if abc[n] not in abc[:n]]
                abc_dict[k_abc] = abc
            prop_merge_dict = dict(ass_dict.items() + abc_dict.items())
            prop_dict[i.name()] = prop_merge_dict
        else:
            for k_ass, v_ass in assembly_prop_dict.items():
                ass_dict = {}
                ass = [x for x in v_ass if x.startswith(i.split(':')[0])]
                ass = [ass[n] for n in range(len(ass)) if ass[n] not in ass[:n]]
                ass_dict[k_ass] = ass
            for k_abc, v_abc in alembic_prop_dict.items():
                abc_dict = {}
                abc = [x for x in v_abc if x.startswith(i.split(':')[0])]
                abc = [abc[n] for n in range(len(abc)) if abc[n] not in abc[:n]]
                abc_dict[k_abc] = abc
            prop_merge_dict = dict(ass_dict.items() + abc_dict.items())
            prop_dict[i.name()].append(prop_merge_dict)

    u"""最后返回一个列表 其中list_dictionary[0]为角色字典， list_dictionary[1]为场景字典, list_dictionary[2]为道具字典"""
    return [chr_dict, bg_dict, prop_dict]
